#!/usr/bin/env python3
"""
AI辅助翻译脚本 - 使用OpenAI兼容API批量翻译summary内容
支持断点续传、错误重试和自定义API端点
支持OpenAI、Azure OpenAI、Claude、国内大模型等多种API服务
"""

import os
import json
import time
import logging
import signal
from typing import Dict, List, Tuple, Optional
import openai
from openai import OpenAI

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class TranslationProgress:
    """翻译进度管理类"""

    def __init__(self, progress_file: str):
        self.progress_file = progress_file
        self.data = {
            'total_records': 0,
            'completed_records': 0,
            'failed_records': 0,
            'current_index': 0,
            'batch_size': 5,
            'start_time': None,
            'last_save_time': None,
            'status': 'pending',  # pending, running, paused, completed, failed
            'completed_ids': [],
            'failed_ids': [],
            'current_batch': []
        }

    def load(self):
        """加载进度数据"""
        if os.path.exists(self.progress_file):
            try:
                with open(self.progress_file, 'r', encoding='utf-8') as f:
                    self.data = json.load(f)
                logger.info(f"加载进度数据: {self.data['completed_records']}/{self.data['total_records']} 已完成")
                return True
            except Exception as e:
                logger.warning(f"加载进度文件失败: {e}")
        return False

    def save(self):
        """保存进度数据"""
        try:
            self.data['last_save_time'] = time.time()
            with open(self.progress_file, 'w', encoding='utf-8') as f:
                json.dump(self.data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.error(f"保存进度文件失败: {e}")

    def initialize(self, total_records: int, batch_size: int = 5):
        """初始化进度数据"""
        self.data.update({
            'total_records': total_records,
            'completed_records': 0,
            'failed_records': 0,
            'current_index': 0,
            'batch_size': batch_size,
            'start_time': time.time(),
            'status': 'running',
            'completed_ids': [],
            'failed_ids': [],
            'current_batch': []
        })
        self.save()

    def mark_completed(self, record_id: int):
        """标记记录为已完成"""
        if record_id not in self.data['completed_ids']:
            self.data['completed_ids'].append(record_id)
            self.data['completed_records'] = len(self.data['completed_ids'])
            self.data['current_index'] += 1
            self.save()

    def mark_failed(self, record_id: int):
        """标记记录为失败"""
        if record_id not in self.data['failed_ids']:
            self.data['failed_ids'].append(record_id)
            self.data['failed_records'] = len(self.data['failed_ids'])
            self.save()

    def set_current_batch(self, batch: List[Tuple[int, str]]):
        """设置当前批次"""
        self.data['current_batch'] = [(rid, text[:100] + '...' if len(text) > 100 else text)
                                      for rid, text in batch]
        self.save()

    def get_progress_info(self) -> str:
        """获取进度信息"""
        if self.data['total_records'] == 0:
            return "未开始"

        completed = self.data['completed_records']
        total = self.data['total_records']
        percentage = (completed / total) * 100 if total > 0 else 0

        elapsed = time.time() - self.data.get('start_time', time.time())
        if completed > 0:
            eta = (elapsed / completed) * (total - completed)
            eta_str = f", 预计剩余: {self._format_time(eta)}"
        else:
            eta_str = ""

        return f"进度: {completed}/{total} ({percentage:.1f}%){eta_str}"

    def _format_time(self, seconds: float) -> str:
        """格式化时间显示"""
        if seconds < 60:
            return f"{seconds:.0f}秒"
        elif seconds < 3600:
            return f"{seconds/60:.1f}分钟"
        else:
            return f"{seconds/3600:.1f}小时"

    def set_status(self, status: str):
        """设置状态"""
        self.data['status'] = status
        self.save()

    def is_completed(self) -> bool:
        """检查是否已完成"""
        return self.data['status'] == 'completed'

    def can_resume(self) -> bool:
        """检查是否可以恢复"""
        return self.data['status'] in ['running', 'paused'] and self.data['completed_records'] < self.data['total_records']

class AITranslator:
    """AI翻译器 - 支持多种OpenAI兼容API"""

    def __init__(self, api_key: str = None, base_url: str = None, model: str = "gpt-3.5-turbo"):
        """
        初始化AI翻译器

        Args:
            api_key: API密钥
            base_url: API基础URL，支持多种服务
            model: 使用的模型名称
        """
        self.api_key = api_key
        self.base_url = base_url
        self.model = model
        self.client = None
        self.translated_count = 0
        self.failed_count = 0
        self.interrupted = False
        self.current_progress = None

        if api_key:
            self._initialize_client()

    def setup_interrupt_handler(self):
        """设置中断信号处理器"""
        def signal_handler(signum, frame):
            logger.info("接收到中断信号，正在安全停止翻译...")
            self.interrupted = True
            if self.current_progress:
                self.current_progress.set_status('paused')
                logger.info(f"进度已保存，可使用 --resume 恢复翻译")
                logger.info(f"当前进度: {self.current_progress.get_progress_info()}")

        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGTERM, signal_handler)

    def _initialize_client(self):
        """初始化API客户端"""
        try:
            client_config = {
                'api_key': self.api_key
            }

            # 如果指定了base_url，添加到配置中
            if self.base_url:
                client_config['base_url'] = self.base_url
                logger.info(f"使用自定义API端点: {self.base_url}")
            else:
                logger.info("使用默认OpenAI API端点")

            self.client = OpenAI(**client_config)
            logger.info(f"API客户端初始化成功，使用模型: {self.model}")

        except Exception as e:
            logger.error(f"API客户端初始化失败: {e}")
            self.client = None

    def set_api_config(self, api_key: str, base_url: str = None, model: str = None):
        """
        设置API配置

        Args:
            api_key: API密钥
            base_url: API基础URL
            model: 模型名称
        """
        self.api_key = api_key
        if base_url:
            self.base_url = base_url
        if model:
            self.model = model

        self._initialize_client()

    def get_supported_providers(self):
        """获取支持的API提供商配置示例"""
        providers = {
            "OpenAI官方": {
                "base_url": None,  # 使用默认
                "model": "gpt-3.5-turbo",
                "env_key": "OPENAI_API_KEY",
                "description": "OpenAI官方API服务"
            },
            "Azure OpenAI": {
                "base_url": "https://your-resource.openai.azure.com/",
                "model": "gpt-35-turbo",  # Azure模型名格式不同
                "env_key": "AZURE_OPENAI_API_KEY",
                "description": "微软Azure OpenAI服务"
            },
            "智谱AI (GLM)": {
                "base_url": "https://open.bigmodel.cn/api/paas/v4/",
                "model": "glm-4",
                "env_key": "ZHIPUAI_API_KEY",
                "description": "智谱AI GLM-4模型"
            },
            "百度千帆": {
                "base_url": "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop/chat/",
                "model": "ernie-3.5-8k",
                "env_key": "BAIDU_API_KEY",
                "description": "百度文心一言"
            },
            "阿里通义": {
                "base_url": "https://dashscope.aliyuncs.com/compatible-mode/v1/",
                "model": "qwen-turbo",
                "env_key": "DASHSCOPE_API_KEY",
                "description": "阿里通义千问"
            },
            "深度求索": {
                "base_url": "https://api.deepseek.com/",
                "model": "deepseek-chat",
                "env_key": "DEEPSEEK_API_KEY",
                "description": "深度求索DeepSeek"
            },
            "月之暗面": {
                "base_url": "https://api.moonshot.cn/v1",
                "model": "moonshot-v1-8k",
                "env_key": "MOONSHOT_API_KEY",
                "description": "月之暗面Kimi"
            },
            "零一万物": {
                "base_url": "https://api.lingyiwanwu.com/v1",
                "model": "yi-34b-chat",
                "env_key": "YI_API_KEY",
                "description": "零一万物Yi系列"
            },
            "Minimax": {
                "base_url": "https://api.minimax.chat/v1",
                "model": "abab6.5-chat",
                "env_key": "MINIMAX_API_KEY",
                "description": "Minimax ABAB系列"
            }
        }
        return providers

    def print_supported_providers(self):
        """打印支持的API提供商信息"""
        providers = self.get_supported_providers()

        print("\n=== 支持的API提供商 ===")
        for name, config in providers.items():
            print(f"\n📡 {name}")
            print(f"   描述: {config['description']}")
            print(f"   模型: {config['model']}")
            print(f"   环境变量: {config['env_key']}")
            if config['base_url']:
                print(f"   端点: {config['base_url']}")
            else:
                print(f"   端点: 使用默认")

        print("\n=== 使用示例 ===")
        print("1. OpenAI官方:")
        print("   export OPENAI_API_KEY='your-key'")
        print("   python ai_translate.py")
        print("\n2. 智谱AI:")
        print("   export AI_API_KEY='your-key'")
        print("   export AI_BASE_URL='https://open.bigmodel.cn/api/paas/v4/'")
        print("   export AI_MODEL='glm-4'")
        print("   python ai_translate.py")

    def test_connection(self):
        """测试API连接"""
        if not self.client:
            logger.error("API客户端未初始化")
            return False

        try:
            logger.info("测试API连接...")
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "user", "content": "测试连接，请回复'连接成功'"}
                ],
                max_tokens=10
            )

            result = response.choices[0].message.content.strip()
            logger.info(f"API连接测试成功，响应: {result}")
            return True

        except Exception as e:
            logger.error(f"API连接测试失败: {e}")
            return False

    def translate_text(self, text: str, context: str = "") -> str:
        """使用AI翻译单条文本"""
        if not text.strip():
            return text

        # 检查是否包含中文
        if self._contains_chinese(text):
            return text

        prompt = f"""请将以下英文技术文档翻译为中文，要求：

1. 保持技术术语的准确性
2. 确保翻译流畅自然
3. 保持专业性和一致性
4. 保留大括号中的占位符如{{CurrentState}}、{{CallType}}等不变，不要翻译这些占位符

上下文：{context}

原文：
{text}

请只返回翻译后的中文，不要添加其他说明。"""

        try:
            response = self.client.chat.completions.create(
                model=self.model,
                messages=[
                    {"role": "system", "content": "你是一个专业的技术文档翻译专家，擅长将英文技术文档准确翻译为中文。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.1,
                max_tokens=1000
            )

            translated = response.choices[0].message.content.strip()
            return translated

        except Exception as e:
            logger.error(f"AI翻译失败: {e}")
            self.failed_count += 1
            return text

    def _contains_chinese(self, text: str) -> bool:
        """检查文本是否包含中文"""
        for char in text:
            if '\u4e00' <= char <= '\u9fff':
                return True
        return False

    def translate_file_batch(self, input_file: str, output_file: str, batch_size: int = 10, resume: bool = False):
        """批量翻译文件，支持中断恢复"""
        logger.info(f"开始批量翻译文件: {input_file}")

        # 设置进度文件路径
        progress_file = input_file.replace('.txt', '_progress.json')
        self.current_progress = TranslationProgress(progress_file)

        # 加载或初始化进度
        if resume and self.current_progress.load():
            logger.info("从之前的进度继续翻译")
            if not self.current_progress.can_resume():
                logger.warning("无法恢复翻译，开始新的翻译任务")
                self.current_progress = TranslationProgress(progress_file)
        else:
            logger.info("开始新的翻译任务")

        # 设置中断处理
        self.setup_interrupt_handler()

        # 读取输入文件
        translations = {}
        try:
            with open(input_file, 'r', encoding='utf-8') as f:
                lines = f.readlines()

            for line in lines:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue

                if '|' in line:
                    parts = line.split('|', 1)
                    try:
                        record_id = int(parts[0].strip())
                        original_text = parts[1].strip()
                        translations[record_id] = original_text
                    except (ValueError, IndexError):
                        continue

            logger.info(f"从文件中读取了 {len(translations)} 条记录")

        except Exception as e:
            logger.error(f"读取输入文件失败: {e}")
            return

        # 初始化进度（如果是新任务）
        if not self.current_progress.load():
            self.current_progress.initialize(len(translations), batch_size)

        # 检查已完成的翻译
        existing_translations = {}
        completed_ids = set(self.current_progress.data['completed_ids'])
        failed_ids = set(self.current_progress.data['failed_ids'])

        # 从输出文件读取已有翻译
        if os.path.exists(output_file):
            try:
                with open(output_file, 'r', encoding='utf-8') as f:
                    lines = f.readlines()
                for line in lines:
                    line = line.strip()
                    if not line or line.startswith('#'):
                        continue
                    if '|' in line:
                        parts = line.split('|', 1)
                        try:
                            record_id = int(parts[0].strip())
                            translated_text = parts[1].strip()
                            existing_translations[record_id] = translated_text
                        except (ValueError, IndexError):
                            continue
                logger.info(f"从已有文件中读取了 {len(existing_translations)} 条翻译")
            except Exception as e:
                logger.warning(f"读取已有翻译文件失败: {e}")

        # 筛选需要翻译的内容
        to_translate = []
        for record_id, original_text in translations.items():
            if record_id not in completed_ids and record_id not in existing_translations:
                to_translate.append((record_id, original_text))

        logger.info(f"需要翻译 {len(to_translate)} 条新记录")
        logger.info(f"已完成: {len(completed_ids)}, 失败: {len(failed_ids)}")

        if not to_translate:
            logger.info("没有需要翻译的内容")
            self.current_progress.set_status('completed')
            logger.info("翻译任务完成！")
            return

        # 批量翻译
        translated_results = existing_translations.copy()
        start_index = self.current_progress.data['current_index']

        logger.info(f"从第 {start_index} 条记录开始翻译")
        logger.info(f"当前进度: {self.current_progress.get_progress_info()}")

        for i in range(start_index, len(to_translate), batch_size):
            # 检查是否被中断
            if self.interrupted:
                logger.info("翻译过程被中断，进度已保存")
                break

            batch = to_translate[i:i + batch_size]
            self.current_progress.set_current_batch(batch)

            logger.info(f"翻译批次: {i+1}-{min(i+batch_size, len(to_translate))}/{len(to_translate)}")
            logger.info(f"当前进度: {self.current_progress.get_progress_info()}")

            for record_id, text in batch:
                # 检查是否被中断
                if self.interrupted:
                    break

                try:
                    logger.info(f"正在翻译记录 {record_id}...")
                    translated = self.translate_text(text)
                    translated_results[record_id] = translated
                    self.translated_count += 1
                    self.current_progress.mark_completed(record_id)

                    # 立即写入文件，支持断点续传
                    self._write_translations(translated_results, output_file)

                except Exception as e:
                    logger.error(f"翻译记录 {record_id} 失败: {e}")
                    self.failed_count += 1
                    self.current_progress.mark_failed(record_id)

                # 添加延迟避免API限流
                if not self.interrupted:
                    time.sleep(0.5)

            # 在批次之间添加更长的延迟
            if i + batch_size < len(to_translate) and not self.interrupted:
                time.sleep(2)

        # 更新最终状态
        if self.interrupted:
            self.current_progress.set_status('paused')
            logger.info(f"翻译已暂停，进度: {self.current_progress.get_progress_info()}")
            logger.info("使用 --resume 参数可以继续翻译")
        else:
            self.current_progress.set_status('completed')
            logger.info("翻译任务完成！")
            logger.info(f"最终统计 - 成功: {self.translated_count}, 失败: {self.failed_count}")

    def resume_translation(self, input_file: str, output_file: str, batch_size: int = 10):
        """恢复中断的翻译任务"""
        logger.info("尝试恢复翻译任务...")
        self.translate_file_batch(input_file, output_file, batch_size, resume=True)

    def show_progress_status(self, input_file: str):
        """显示翻译进度状态"""
        progress_file = input_file.replace('.txt', '_progress.json')
        progress = TranslationProgress(progress_file)

        if not progress.load():
            logger.info("没有找到进度文件，该文件可能还没有开始翻译")
            return

        print("\n=== 翻译进度状态 ===")
        print(f"状态: {progress.data['status']}")
        print(f"总记录数: {progress.data['total_records']}")
        print(f"已完成: {progress.data['completed_records']}")
        print(f"失败: {progress.data['failed_records']}")
        print(f"进度: {progress.get_progress_info()}")

        if progress.can_resume():
            print("\n✅ 可以继续翻译 (使用 --resume)")
        elif progress.is_completed():
            print("\n✅ 翻译已完成")
        else:
            print(f"\n⚠️  状态: {progress.data['status']}")

        if progress.data['start_time']:
            elapsed = time.time() - progress.data['start_time']
            print(f"已用时间: {progress._format_time(elapsed)}")

    def _write_translations(self, translations: Dict[int, str], output_file: str):
        """写入翻译结果到文件"""
        try:
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(f"# AI翻译结果文件\n")
                f.write(f"# 总记录数: {len(translations)}\n")
                f.write("# 格式: ID|Translated Summary\n")
                f.write("\n")

                # 按ID排序
                sorted_items = sorted(translations.items())

                for record_id, translated_text in sorted_items:
                    # 转义管道符号
                    escaped_text = translated_text.replace('|', '\\|')
                    f.write(f"{record_id}|{escaped_text}\n")

        except Exception as e:
            logger.error(f"写入翻译文件失败: {e}")

def main():
    """主函数"""
    import argparse

    # 解析命令行参数
    parser = argparse.ArgumentParser(description='AI翻译脚本 - 支持多种OpenAI兼容API，支持中断恢复')
    parser.add_argument('--test', action='store_true', help='测试API连接')
    parser.add_argument('--list-providers', action='store_true', help='列出支持的API提供商')
    parser.add_argument('--api-key', help='API密钥')
    parser.add_argument('--base-url', help='API基础URL')
    parser.add_argument('--model', help='使用的模型名称')
    parser.add_argument('--batch-size', type=int, default=5, help='批处理大小')
    parser.add_argument('--files', nargs='+', help='指定要翻译的文件')
    parser.add_argument('--resume', action='store_true', help='恢复中断的翻译任务')
    parser.add_argument('--status', help='查看指定文件的翻译状态')
    parser.add_argument('--clean-progress', action='store_true', help='清理进度文件，重新开始翻译')

    args = parser.parse_args()

    # 创建翻译器实例
    translator = AITranslator()

    # 处理特殊命令
    if args.list_providers:
        translator.print_supported_providers()
        return

    # 处理状态查看命令
    if args.status:
        if os.path.exists(args.status):
            translator.show_progress_status(args.status)
        else:
            logger.error(f"文件不存在: {args.status}")
        return

    # 处理清理进度命令
    if args.clean_progress:
        if args.files:
            for file_pattern in args.files:
                if file_pattern in ['metrics', 'variants', 'enum']:
                    if file_pattern == 'metrics':
                        input_file = 'metrics_summaries.txt'
                    elif file_pattern == 'variants':
                        input_file = 'variants_summaries.txt'
                    elif file_pattern == 'enum':
                        input_file = 'enum_summaries.txt'
                else:
                    input_file = file_pattern

                progress_file = input_file.replace('.txt', '_progress.json')
                if os.path.exists(progress_file):
                    os.remove(progress_file)
                    logger.info(f"已清理进度文件: {progress_file}")
                else:
                    logger.info(f"进度文件不存在: {progress_file}")
        else:
            logger.error("请指定要清理进度文件的文件")
        return

    # 获取API配置，优先级：命令行参数 > 环境变量 > 默认值
    api_key = (args.api_key or
               os.getenv('AI_API_KEY') or
               os.getenv('OPENAI_API_KEY') or
               os.getenv('AZURE_OPENAI_API_KEY'))

    base_url = (args.base_url or
                os.getenv('AI_BASE_URL') or
                os.getenv('OPENAI_BASE_URL'))

    model = (args.model or
             os.getenv('AI_MODEL') or
             os.getenv('OPENAI_MODEL', 'gpt-3.5-turbo'))

    if not api_key:
        logger.error("未找到API密钥！")
        logger.error("\n请通过以下方式之一设置API密钥：")
        logger.error("1. 命令行参数: --api-key 'your-key'")
        logger.error("2. 环境变量: export AI_API_KEY='your-key'")
        logger.error("3. OpenAI: export OPENAI_API_KEY='your-key'")
        logger.error("4. 查看 --list-providers 获取更多选项")
        return

    # 配置翻译器
    translator.set_api_config(api_key, base_url, model)

    # 测试连接
    if args.test:
        if translator.test_connection():
            logger.info("API连接测试成功！")
        else:
            logger.error("API连接测试失败！")
        return

    # 默认翻译文件列表
    default_files = [
        ('metrics_summaries.txt', 'metrics_summaries_zh.txt'),
        ('variants_summaries.txt', 'variants_summaries_zh.txt'),
        ('enum_summaries.txt', 'enum_summaries_zh.txt')
    ]

    # 处理用户指定的文件
    if args.files:
        translation_tasks = []
        for file_pattern in args.files:
            if file_pattern in ['metrics', 'variants', 'enum']:
                if file_pattern == 'metrics':
                    translation_tasks.append(('metrics_summaries.txt', 'metrics_summaries_zh.txt'))
                elif file_pattern == 'variants':
                    translation_tasks.append(('variants_summaries.txt', 'variants_summaries_zh.txt'))
                elif file_pattern == 'enum':
                    translation_tasks.append(('enum_summaries.txt', 'enum_summaries_zh.txt'))
            else:
                # 如果用户直接提供了文件路径
                input_file = file_pattern
                output_file = file_pattern.replace('.txt', '_zh.txt')
                translation_tasks.append((input_file, output_file))
    else:
        translation_tasks = default_files

    logger.info(f"开始翻译任务，使用API: {base_url or 'OpenAI官方'}")
    logger.info(f"使用模型: {model}")
    logger.info(f"批处理大小: {args.batch_size}")

    # 执行翻译任务
    for input_file, output_file in translation_tasks:
        if not os.path.exists(input_file):
            logger.warning(f"输入文件不存在，跳过: {input_file}")
            continue

        try:
            logger.info(f"开始翻译: {input_file} → {output_file}")

            if args.resume:
                logger.info("恢复中断的翻译任务")
                translator.resume_translation(input_file, output_file, args.batch_size)
            else:
                translator.translate_file_batch(input_file, output_file, args.batch_size, resume=False)

        except KeyboardInterrupt:
            logger.info("用户中断翻译过程")
            logger.info("使用 --resume 可以继续翻译")
            break
        except Exception as e:
            logger.error(f"翻译文件 {input_file} 失败: {e}")
            continue

    logger.info("所有翻译任务完成！")
    logger.info(f"成功翻译: {translator.translated_count} 条，失败: {translator.failed_count} 条")

if __name__ == "__main__":
    main()